#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// TODO: 翻译你优化后的IR为RISC-V目标代码
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明外部函数 */
    codegen->append_inst(".extern getint", ASMInstruction::Atrribute);
    codegen->append_inst(".extern putint", ASMInstruction::Atrribute);

    /* 声明 main 函数 */ 
    codegen->append_inst(".align 2", ASMInstruction::Atrribute);
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label);

    /* main 函数的 Prologue (序言) */
    codegen->append_inst("addi", {"sp", "sp", "-64"});
    codegen->append_inst("sd", {"ra", "sp", "56"});
    codegen->append_inst("sd", {"s0", "sp", "48"});
    codegen->append_inst("addi", {"s0", "sp", "64"});

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);
    codegen->append_inst("# %op4 = call i32 @getint()", ASMInstruction::Comment);
    codegen->append_inst("call", {"getint"});
    codegen->append_inst("mv", {"s1", "a0"});  // 保存getint的结果到s1 (c的值)
    codegen->append_inst("# br label %label5", ASMInstruction::Comment);
    codegen->append_inst("li", {"t0", "1"});   // 初始化a的值为1
    codegen->append_inst("li", {"t1", "1"});   // 初始化b的值为1
    codegen->append_inst("j", {".main_label5"});

    /* label5: 主循环 */
    codegen->append_inst(".main_label5", ASMInstruction::Label);
    codegen->append_inst("# %op6 = phi i32 [ 1, %label_entry ], [ %op31, %label26 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op_b = phi i32 [ 1, %label_entry ], [ %op_b_final, %label26 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op7 = icmp slt i32 %op6, 99999999", ASMInstruction::Comment);
    codegen->append_inst("li", {"t2", "99999999"});
    codegen->append_inst("slt", {"t3", "t0", "t2"});
    codegen->append_inst("# br i1 %op7, label %label8, label %label15", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t4", "zero", "0"});
    codegen->append_inst("blt", {"t4", "t3", ".main_label8"});
    codegen->append_inst("j", {".main_label15"});

    /* label8: 第一个条件检查 (a+1399)%c==0 */
    codegen->append_inst(".main_label8", ASMInstruction::Label);
    codegen->append_inst("# %op11 = add i32 %op6, 1399", ASMInstruction::Comment);
    codegen->append_inst("li", {"t2", "1399"});
    codegen->append_inst("add", {"t3", "t0", "t2"});  // t3 = a + 1399
    codegen->append_inst("# %op13 = srem i32 %op11, %op4", ASMInstruction::Comment);
    codegen->append_inst("rem", {"t4", "t3", "s1"});  // t4 = (a+1399) % c
    codegen->append_inst("# %op14 = icmp eq i32 %op13, 0", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t5", "zero", "0"});
    codegen->append_inst("# br i1 %op14, label %label18, label %label19", ASMInstruction::Comment);
    codegen->append_inst("beq", {"t4", "t5", ".main_label18"});
    codegen->append_inst("j", {".main_label19"});

    /* label15: 循环结束 */
    codegen->append_inst(".main_label15", ASMInstruction::Label);
    codegen->append_inst("# call void @putint(i32 %op4)", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "s1"});
    codegen->append_inst("call", {"putint"});
    codegen->append_inst("# ret i32 %op6", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "t0"});
    codegen->append_inst("j", {".main_exit"});

    /* label18: (a+v1)%c==0的情况, b=411 */
    codegen->append_inst(".main_label18", ASMInstruction::Label);
    codegen->append_inst("li", {"t1", "411"});
    codegen->append_inst("# br label %label26", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label26"});

    /* label19: 第二个条件检查 (a+1399)%c==1 */
    codegen->append_inst(".main_label19", ASMInstruction::Label);
    codegen->append_inst("# %op22 = add i32 %op6, 1399", ASMInstruction::Comment);
    codegen->append_inst("li", {"t2", "1399"});
    codegen->append_inst("add", {"t3", "t0", "t2"});  // t3 = a + 1399
    codegen->append_inst("# %op24 = srem i32 %op22, %op4", ASMInstruction::Comment);
    codegen->append_inst("rem", {"t4", "t3", "s1"});  // t4 = (a+1399) % c
    codegen->append_inst("# %op25 = icmp eq i32 %op24, 1", ASMInstruction::Comment);
    codegen->append_inst("li", {"t5", "1"});
    codegen->append_inst("# br i1 %op25, label %label32, label %label33", ASMInstruction::Comment);
    codegen->append_inst("beq", {"t4", "t5", ".main_label32"});
    codegen->append_inst("j", {".main_label33"});

    /* label26: 计算最终结果并返回循环 */
    codegen->append_inst(".main_label26", ASMInstruction::Label);
    codegen->append_inst("# %op_b_final = phi i32 [ 411, %label18 ], [ %op_b_from_40, %label40 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op28 = add i32 %op6, 1", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t2", "t0", "1"});  // t2 = a + 1
    codegen->append_inst("# %op30 = srem i32 %op_b_final, 3", ASMInstruction::Comment);
    codegen->append_inst("li", {"t3", "3"});
    codegen->append_inst("rem", {"t4", "t1", "t3"});  // t4 = b % 3
    codegen->append_inst("# %op31 = add i32 %op28, %op30", ASMInstruction::Comment);
    codegen->append_inst("add", {"t0", "t2", "t4"});  // a = (a + 1) + (b % 3)
    codegen->append_inst("# br label %label5", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label5"});

    /* label32: (a+v1)%c==1的情况, b=101 */
    codegen->append_inst(".main_label32", ASMInstruction::Label);
    codegen->append_inst("li", {"t1", "101"});
    codegen->append_inst("# br label %label40", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label40"});

    /* label33: 第三个条件检查 (a+1399)%c==3 */
    codegen->append_inst(".main_label33", ASMInstruction::Label);
    codegen->append_inst("# %op36 = add i32 %op6, 1399", ASMInstruction::Comment);
    codegen->append_inst("li", {"t2", "1399"});
    codegen->append_inst("add", {"t3", "t0", "t2"});  // t3 = a + 1399
    codegen->append_inst("# %op38 = srem i32 %op36, %op4", ASMInstruction::Comment);
    codegen->append_inst("rem", {"t4", "t3", "s1"});  // t4 = (a+1399) % c
    codegen->append_inst("# %op39 = icmp eq i32 %op38, 3", ASMInstruction::Comment);
    codegen->append_inst("li", {"t5", "3"});
    codegen->append_inst("# br i1 %op39, label %label41, label %label42", ASMInstruction::Comment);
    codegen->append_inst("beq", {"t4", "t5", ".main_label41"});
    codegen->append_inst("j", {".main_label42"});

    /* label40: 从label32来的分支合并点 */
    codegen->append_inst(".main_label40", ASMInstruction::Label);
    codegen->append_inst("# %op_b_from_40 = phi i32 [ 101, %label32 ], [ %op_b_from_42, %label42 ]", ASMInstruction::Comment);
    codegen->append_inst("# br label %label26", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label26"});

    /* label41: (a+v1)%c==3的情况, b=99 */
    codegen->append_inst(".main_label41", ASMInstruction::Label);
    codegen->append_inst("li", {"t1", "99"});
    codegen->append_inst("# br label %label42", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label42"});

    /* label42: 最后的分支合并点 */
    codegen->append_inst(".main_label42", ASMInstruction::Label);
    codegen->append_inst("# %op_b_from_42 = phi i32 [ 99, %label41 ], [ %op_b, %label33 ]", ASMInstruction::Comment);
    codegen->append_inst("# br label %label40", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label40"});

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst(".main_exit", ASMInstruction::Label);
    codegen->append_inst("ld", {"ra", "sp", "56"});
    codegen->append_inst("ld", {"s0", "sp", "48"});
    codegen->append_inst("addi", {"sp", "sp", "64"});
    codegen->append_inst("ret");
}
